from utils.page_util import PageResponseModel


def merge_associated_data(data_list, association_mapping=None, default_value=''):
    """
    合并关联数据

    :param data_list: 主数据列表，格式为 [{}, ..., {}] 或 [[{}, {}, ...], ..., [{}, {}, ...]]
    :param association_mapping: 关联映射配置，格式为 {'target_field': source_index}
                               例如: {'stockName': 1, 'accountName': 2}
    :param default_value: 默认值
    :return: 合并后的数据列表
    """
    if not data_list:
        return data_list

    if association_mapping is None:
        association_mapping = {'stockName': 1}

    merged_list = []
    for item in data_list:
        # 如果是元组或列表，说明是关联查询的结果
        if isinstance(item, (list, tuple)):
            # 第一个元素是主表数据
            if isinstance(item[0], dict):
                merged_item = item[0].copy()
            else:
                merged_item = dict(item[0])

            # 根据映射关系合并关联数据
            for target_field, source_index in association_mapping.items():
                if source_index < len(item):
                    associated_data = item[source_index]
                    merged_item[target_field] = associated_data if associated_data is not None else default_value
                else:
                    merged_item[target_field] = default_value
        else:
            # 如果不是元组或列表，直接使用
            merged_item = item

        merged_list.append(merged_item)

    return merged_list


def merge_page_associated_data(page_result, association_mapping=None, default_value=''):
    """
    合并分页数据中的关联数据

    :param page_result: 分页结果对象
    :param association_mapping: 关联映射配置，格式为 {'target_field': source_index}
    :param default_value: 默认值
    :return: 处理后的分页结果对象
    """
    if not page_result or not page_result.rows:
        return page_result

    if association_mapping is None:
        association_mapping = {'stockName': 1}

    merged_rows = merge_associated_data(page_result.rows, association_mapping, default_value)

    return PageResponseModel(**{**page_result.model_dump(by_alias=True), 'rows': merged_rows})


def process_associated_data(data, association_mapping=None, default_value='', is_page=False):
    """
    统一处理关联数据的入口方法

    :param data: 数据，可以是分页结果对象或普通列表
    :param association_mapping: 关联映射配置
    :param default_value: 默认值
    :param is_page: 是否为分页数据
    :return: 处理后的数据
    """
    if is_page:
        return merge_page_associated_data(data, association_mapping, default_value)
    else:
        return merge_associated_data(data, association_mapping, default_value)
